/*
 *  fuzz.mm
 *  bbinj
 *
 *  Created by msftguy on 6/20/10.
 *  Copyright 2010 __MyCompanyName__. All rights reserved.
 *
 */

#include "fuzz.h"

#include "logging.h"

#include <mach/mach.h>
#include <mach-o/dyld.h>
#include <mach-o/getsect.h>

char* findBytesInSect(const mach_header* mh, size_t slide, const char* segname, const char* sectname, const char* pattern, 
					  size_t patternSize = 0, int align = 0);

char* findBytesInSect(const mach_header* mh, size_t slide, const char* segname, const char* sectname, const char* pattern, 
					  size_t patternSize, int align)
{
	if (patternSize == 0) {
		patternSize = strlen(pattern) + 1;
	}
	if (align == 0) {
		align = 1;
	}
	uint32_t sectSize = 0;
	char* sectData = getsectdatafromheader(mh, segname, sectname, &sectSize);
	if (sectData == NULL) {
		log_progress("findBytesInSect: section %s:%s not found", segname, sectname);	
	}
	sectData += slide;

	char* foundLoc = NULL;
	for (char* p = sectData; p < sectData + sectSize - patternSize; p += align) 
	{
		if (memcmp(p, pattern, patternSize) == 0) {
			if (foundLoc == NULL) {
				foundLoc = p;
			} else {
				log_progress("findBytesInSect: multiple matches: %p and %p", foundLoc, p);
				return NULL;
			}
		}
	}
	return foundLoc;
}

bool find_at_handler(void** ppAtHandler, void** pPingResponseXref)
{
	const mach_header* mh = _dyld_get_image_header(0);
	size_t slide = _dyld_get_image_vmaddr_slide(0);
	char* atDispatcherName19 = findBytesInSect(mh, slide, "__TEXT", "__cstring", "19ATDispatcherPrivate");
	if (atDispatcherName19 == NULL) {
		log_progress("find_at_handler: can't find AT command dispatcher class name");
		return false;
	}
	
	log_progress("atDispatcherName19 at %p", atDispatcherName19);
		
	char* atDispatcherNameXref = findBytesInSect(mh, slide, "__DATA", "__const", (char*)&atDispatcherName19, 4, 4);
	if (atDispatcherNameXref == NULL) {
		log_progress("find_at_handler: can't find AT command dispatcher class name xref");
		return false;
	}
	// vtable sanity check
	{
		uint32_t u1 = *(uint32_t*)(atDispatcherNameXref - 0x4C);
		uint32_t u2 = (uint32_t)(atDispatcherNameXref - 4);
		if (u1 != u2) {
			log_progress("find_at_handler: AT command dispatcher VTBL check failed: %u != %u", u1, u2);
			return false;
		}
	}
	
	char* handlePingResponseStr = findBytesInSect(mh, slide, "__TEXT", "__cstring", "handlePingResponse\n");
	if (handlePingResponseStr == NULL) {
		log_progress("find_at_handler: can't find 'handlePingResponse' str");
		return false;		
	}
	
	char* handlePingResponseXref = findBytesInSect(mh, slide, "__TEXT", "__text", (char*)&handlePingResponseStr, 4, 4);
	if (handlePingResponseXref == NULL) {
		log_progress("find_at_handler: can't find 'handlePingResponse' xref");
		return false;		
	}
	
	*pPingResponseXref = handlePingResponseXref;
	
	*ppAtHandler = (void*)(atDispatcherNameXref - 0x20);
	return true;
}

bool patch_word(void* addr, void* replaceWith)
{
	vm_region_basic_info info;
	vm_address_t vmaddr = (vm_address_t) addr;
	vm_size_t vmsize = sizeof(void*);
	mach_msg_type_number_t infoCount = sizeof(info);
	mach_port_t	object_name;
	kern_return_t result = vm_region(mach_task_self(), &vmaddr, &vmsize, VM_REGION_BASIC_INFO, (int*)&info, &infoCount, &object_name);
	if (result != KERN_SUCCESS) {
		log_progress("error getting page protection: err=0x%X", result);
		return false;
	}
	result = vm_protect(mach_task_self(), vmaddr, sizeof(void*), false, VM_PROT_WRITE | VM_PROT_COPY);
	if (result != KERN_SUCCESS) {
		log_progress("error setting page protection: err=0x%X", result);
		return false;
	}
	*(uint32_t*)addr =  (uint32_t)replaceWith;
	result = vm_protect(mach_task_self(), vmaddr, sizeof(void*), false, info.protection);
	if (result != KERN_SUCCESS) {
		log_progress("error restoring page protection: err=0x%X", result);
		return false;
	}
	return true;
}
